import logging
import re
from typing import List

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain_core.language_models import BaseLLM

from langchain.chains import MapReduceDocumentsChain, ReduceDocumentsChain
from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.chains.llm import LLMChain
from langchain_community.llms.ollama import Ollama
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import PromptTemplate
from langchain_core.documents import Document
from langchain_ollama import OllamaEmbeddings

# 添加自定义tokenizer避免连接HuggingFace
def custom_get_tokenizer():
    # 使用简单的空格分词作为替代方案
    return lambda text: text.split()

# 覆盖默认的get_tokenizer函数
BaseLLM.get_token_ids = lambda self, text: custom_get_tokenizer()(text)

# 配置
EMBEDDING_URL="http://192.168.124.23:11434"
EMBEDDING_MODEL = "nomic-embed-text"
CHAT_URL = "http://192.168.124.23:11434"
CHAT_MODEL = 'qwen3:1.7b'

# 初始化模型
try:
    embeddings = OllamaEmbeddings(
        model=EMBEDDING_MODEL,
        base_url=EMBEDDING_URL
    )

    # 加载向量数据库
    db = FAISS.load_local(
        r'./faiss/camp',
        embeddings,
        allow_dangerous_deserialization=True
    )

    chat_model = Ollama(
        model=CHAT_MODEL,
        base_url=CHAT_URL,
        temperature=0.3,  # 降低随机性
        top_p=0.85,  # 提高回答确定性
        num_ctx=8192,  # 增加上下文长度
        repeat_penalty=1.1  # 减少重复内容
    )
except Exception as e:
    logging.error(f"初始化模型失败: {e}")
    raise

async def get_related_content(related_docs):
    related_content=[]
    for doc in related_docs:
        related_content.append(doc.page_content.replace('\n\n', '\n'))
    return '\n'.join(related_content)


async def define_prompt(question:str):
    docs=db.similarity_search(question, k=10, fetch_k=40)
    logging.info(f"相关文档: {docs}")
    related_docs = await get_related_content(docs)

    PROMPT_TEMPLATE="""
        基于以下已知信息，简洁和专业地回答用户的问题。不允许在回答中添加任何编造成分。
        已知内容：
        {context}
        问题：
        {question}
    """
    prompt = PromptTemplate(input_variables=["context","question"], template=PROMPT_TEMPLATE)
    return prompt.format(context=related_docs, question=question)

async def qa(question):
    llm=chat_model
    my_prompt= await define_prompt(question)
    result=llm.invoke(my_prompt)
    result=re.sub(r"(<think>.*?</think>\s+)","",result,flags=re.S)
    return result



async def rag_query_by_file(question: str) -> str:
    """使用Map-Reduce方式的RAG模型查询问题"""
    try:
        # 检索相关文档
        # 混合检索策略
        base_retriever = db.as_retriever(
            search_type="mmr",
            search_kwargs={
                'k': 10,
                'fetch_k': 40,
                'lambda_mult': 0.4  # 多样性参数
            }
        )

        # 检索相关文档
        retriever_chunks = base_retriever.invoke(question)
        logging.info(f"检索到 {len(retriever_chunks)} 个相关文档块\n")
        logging.info(f"检索到的相关内容: {retriever_chunks}")

        # 添加后段内容增强策略
        # 识别可能遗漏的后段关键词
        if len(retriever_chunks) > 0:
            # 提取前段内容中的关键词
            context_keywords = extract_keywords(" ".join([doc.page_content[:500] for doc in retriever_chunks]))

            # 检索可能遗漏的后段内容
            if context_keywords:
                supplementary_chunks = base_retriever.invoke(" ".join(context_keywords))
                # 合并结果，过滤重复
                all_chunks = retriever_chunks + [
                    doc for doc in supplementary_chunks
                    if doc not in retriever_chunks
                ]
                retriever_chunks = all_chunks[:15]  # 限制总数

        logging.info(f"增强后共 {len(retriever_chunks)} 个文档块")

        # Map阶段：处理每个文档
        map_template = """请仔细分析以下文档片段，找出与问题直接相关的信息：
        --- 文档片段 ---
        {docs}
        --- 结束 ---
        
        问题: {question}
        
        提取要求:
        1. 特别关注片段中的关键事实和数据
        2. 如果涉及流程或顺序，注意片段中的步骤描述
        3. 如果文档中没有相关信息，回答"无相关内容"
        4. 保持回答简洁，不超过80字
        
        基于文档的回答:"""
        map_prompt = PromptTemplate.from_template(map_template)
        map_chain = LLMChain(llm=chat_model, prompt=map_prompt)

        # Reduce阶段：合并所有回答
        reduce_template = """您是一位信息整合专家，请根据以下多个回答合成准确、完整的最终答案:
        --- 各片段回答 ---
        {doc_summaries}
        --- 结束 ---
        
        原始问题: {question}
        
        整合要求:
        1. 综合所有信息形成完整答案
        2. 对于流程类问题，按正确时间顺序组织
        3. 对于概念类问题，给出清晰定义
        4. 如果存在冲突信息，优先使用最新文档
        5. 如果所有回答都是"无相关内容"，则最终答案也为"无相关信息"
        
        最终答案结构:
        - 以"根据文档内容"开头
        - 按逻辑分段组织内容
        - 结尾添加"（来源：文档综合）"
        
        最终答案:"""
        reduce_prompt = PromptTemplate.from_template(reduce_template)
        reduce_chain = LLMChain(llm=chat_model, prompt=reduce_prompt)

        # 合并文档链
        combine_documents_chain = StuffDocumentsChain(
            llm_chain=reduce_chain,
            document_variable_name="doc_summaries"
        )

        # 减少文档链
        reduce_documents_chain = ReduceDocumentsChain(
            combine_documents_chain=combine_documents_chain,
            collapse_documents_chain=combine_documents_chain,
            token_max=8000,
        )

        # 映射规约链
        map_reduce_chain = MapReduceDocumentsChain(
            llm_chain=map_chain,
            reduce_documents_chain=reduce_documents_chain,
            document_variable_name="docs",
            return_intermediate_steps=False,
        )

        # 执行查询
        result = map_reduce_chain.invoke({
            "input_documents": retriever_chunks,
            "question": question
        })

        logging.info(f"查询结果: {result['output_text']}")
        return result['output_text']

    except Exception as e:
        logging.error(f"查询过程中出错: {e}")
        return f"查询过程中出现错误: {str(e)}"


def extract_keywords(text: str, top_n: int = 5) -> List[str]:
    """提取文本中的关键词"""
    # 简单实现 - 实际中可用更复杂的NLP技术
    words = re.findall(r'\b\w{3,}\b', text.lower())
    word_freq = {}
    for word in words:
        if word not in ["的", "是", "在", "和", "有"]:  # 过滤停用词
            word_freq[word] = word_freq.get(word, 0) + 1

    # 取频率最高的词
    sorted_words = sorted(word_freq.items(), key=lambda x: x[1], reverse=True)
    return [word for word, freq in sorted_words[:top_n]]

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    question = "申请报废机动车回收拆解资质需满足哪些场地要求？"
    try:
        #res = rag_query_by_file(question)
        res=qa(question)
        print(f"查询结果: {res}")
    except Exception as e:
        print(f"程序运行出错: {e}")


